package com.tianji.aigc.advisor;

import cn.hutool.core.convert.Convert;
import com.tianji.aigc.enums.AgentTypeEnum;
import com.tianji.aigc.memory.MyChatMemory;
import org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.api.*;
import org.springframework.ai.chat.model.ChatResponse;
import reactor.core.publisher.Flux;

/**
 * 对话记录优化的增强器
 */
public class RecordOptimizationAdvisor implements CallAroundAdvisor, StreamAroundAdvisor {

    private MyChatMemory myChatMemory;

    public RecordOptimizationAdvisor(MyChatMemory myChatMemory) {
        this.myChatMemory = myChatMemory;
    }

    @Override
    public AdvisedResponse aroundCall(AdvisedRequest advisedRequest, CallAroundAdvisorChain chain) {
        AdvisedResponse advisedResponse = chain.nextAroundCall(advisedRequest);
        ChatResponse response = advisedResponse.response();
        String text = response.getResult().getOutput().getText();
        AgentTypeEnum agentTypeEnum = AgentTypeEnum.agentNameOf(text);
        if (null != agentTypeEnum) {
            // 路由智能体 需要 转发请求到 下游的智能体了，这里就需要优化记录了
            var key = AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY;
            var conversationId = Convert.toStr(advisedRequest.adviseContext().get(key));
            this.myChatMemory.optimization(conversationId);
        }
        return advisedResponse;
    }

    @Override
    public Flux<AdvisedResponse> aroundStream(AdvisedRequest advisedRequest, StreamAroundAdvisorChain chain) {
        // 流式的调用，不做处理
        return chain.nextAroundStream(advisedRequest);
    }

    @Override
    public String getName() {
        return this.getClass().getSimpleName();
    }

    @Override
    public int getOrder() {
        return Advisor.DEFAULT_CHAT_MEMORY_PRECEDENCE_ORDER - 100;
    }
}
